import argparse
import os
import shutil
from datetime import datetime, timedelta


class TestCleaner:
    def __init__(self):
        self.base_dir = os.path.dirname(os.path.abspath(__file__))
        self.reports_dir = os.path.join(self.base_dir, "reports")
        self.logs_dir = os.path.join(self.base_dir, "logs")
    # days 参数清理前几天的数据
    def clean_reports(self, days=7, dry_run=False):
        """清理旧的测试报告"""
        print(f"清理 {days} 天前的测试报告...")
        if dry_run:
            print("=== 干跑模式，不会删除任何文件 ===")

        count = 0
        size = 0

        if not os.path.exists(self.reports_dir):
            print(f"报告目录不存在: {self.reports_dir}")
            return count, size

        now = datetime.now()

        for item in os.listdir(self.reports_dir):
            item_path = os.path.join(self.reports_dir, item)

            # 跳过allure-report目录（最新报告）
            if item == "allure-report":
                continue

            # 获取文件/目录的修改时间
            try:
                mtime = datetime.fromtimestamp(os.path.getmtime(item_path))
            except OSError:
                print(f"无法获取时间戳: {item_path}")
                continue

            # 计算文件/目录的大小
            item_size = self._get_size(item_path)

            # 如果修改时间超过指定天数，则删除
            if now - mtime > timedelta(days=days):
                print(f"将删除: {item_path} ({self._format_size(item_size)})")
                if not dry_run:
                    try:
                        if os.path.isfile(item_path):
                            os.remove(item_path)
                        else:
                            shutil.rmtree(item_path)
                        count += 1
                        size += item_size
                    except Exception as e:
                        print(f"删除失败: {item_path}, 错误: {e}")

        return count, size

    def clean_logs(self, days=7, dry_run=False):
        """清理旧的日志文件"""
        print(f"清理 {days} 天前的日志文件...")
        if dry_run:
            print("=== 干跑模式，不会删除任何文件 ===")

        count = 0
        size = 0

        if not os.path.exists(self.logs_dir):
            print(f"日志目录不存在: {self.logs_dir}")
            return count, size

        now = datetime.now()

        for item in os.listdir(self.logs_dir):
            item_path = os.path.join(self.logs_dir, item)

            # 获取文件的修改时间
            try:
                mtime = datetime.fromtimestamp(os.path.getmtime(item_path))
            except OSError:
                print(f"无法获取时间戳: {item_path}")
                continue

            # 计算文件大小
            item_size = os.path.getsize(item_path)

            # 如果修改时间超过指定天数，则删除
            if now - mtime > timedelta(days=days):
                print(f"将删除: {item_path} ({self._format_size(item_size)})")
                if not dry_run:
                    try:
                        os.remove(item_path)
                        count += 1
                        size += item_size
                    except Exception as e:
                        print(f"删除失败: {item_path}, 错误: {e}")

        return count, size

    def _get_size(self, path):
        """计算文件或目录的大小"""
        if os.path.isfile(path):
            return os.path.getsize(path)

        total_size = 0
        for dirpath, dirnames, filenames in os.walk(path):
            for f in filenames:
                fp = os.path.join(dirpath, f)
                try:
                    total_size += os.path.getsize(fp)
                except OSError:
                    pass
        return total_size

    def _format_size(self, size_bytes):
        """格式化文件大小为人类可读的形式"""
        units = ['B', 'KB', 'MB', 'GB', 'TB']
        unit_index = 0
        while size_bytes >= 1024 and unit_index < len(units) - 1:
            size_bytes /= 1024
            unit_index += 1
        return f"{size_bytes:.2f} {units[unit_index]}"


def main():
    parser = argparse.ArgumentParser(description="清理测试报告和日志工具")
    parser.add_argument("-d", "--days", type=int, default=7,
                        help="保留最近多少天的文件，默认为7天")
    parser.add_argument("--dry-run", action="store_true",
                        help="干跑模式，不实际删除文件，只显示将要删除的内容")
    parser.add_argument("--reports", action="store_true", help="只清理报告")
    parser.add_argument("--logs", action="store_true", help="只清理日志")

    args = parser.parse_args()

    cleaner = TestCleaner()

    if not args.reports or args.logs:
        report_count, report_size = cleaner.clean_reports(args.days, args.dry_run)
    else:
        report_count, report_size = 0, 0

    if not args.logs or args.reports:
        log_count, log_size = cleaner.clean_logs(args.days, args.dry_run)
    else:
        log_count, log_size = 0, 0

    total_count = report_count + log_count
    total_size = report_size + log_size

    if not args.dry_run:
        print(f"清理完成！共删除 {total_count} 个文件/目录，释放空间: {cleaner._format_size(total_size)}")
    else:
        print(f"干跑模式完成！将删除 {total_count} 个文件/目录，释放空间: {cleaner._format_size(total_size)}")


if __name__ == "__main__":
    main()
